import torch.nn as nn
from resnet18_backbone import Resnet18
from PiCANet import GlobalPiCANet, LocalPiCANet,MultiScalePiCANet
import torch.nn.functional as F

class PiCANetResnet18(nn.Module):
    def __init__(self, model_type='full'):
        super().__init__()
        self.backbone = Resnet18()

        # 根据模型类型配置注意力模块
        self.model_type = model_type

        if model_type == 'full':
            self.pica3 = MultiScalePiCANet(256, use_local=True, use_mid=True, use_global=True)
            self.pica4 = MultiScalePiCANet(512, use_local=True, use_mid=True, use_global=True)
        elif model_type == 'only_local':
            self.pica3 = MultiScalePiCANet(256, use_local=True, use_mid=False, use_global=False)
            self.pica4 = MultiScalePiCANet(512, use_local=True, use_mid=False, use_global=False)
        elif model_type == 'only_global':
            self.pica3 = MultiScalePiCANet(256, use_local=False, use_mid=False, use_global=True)
            self.pica4 = MultiScalePiCANet(512, use_local=False, use_mid=False, use_global=True)
        elif model_type == 'only_mid':
            self.pica3 = MultiScalePiCANet(256, use_local=False, use_mid=True, use_global=False)
            self.pica4 = MultiScalePiCANet(512, use_local=False, use_mid=True, use_global=False)
        elif model_type == 'local_global':
            self.pica3 = MultiScalePiCANet(256, use_local=True, use_mid=False, use_global=True)
            self.pica4 = MultiScalePiCANet(512, use_local=True, use_mid=False, use_global=True)
        elif model_type == 'local_mid':
            self.pica3 = MultiScalePiCANet(256, use_local=True, use_mid=True, use_global=False)
            self.pica4 = MultiScalePiCANet(512, use_local=True, use_mid=True, use_global=False)
        elif model_type == 'global_mid':
            self.pica3 = MultiScalePiCANet(256, use_local=False, use_mid=True, use_global=True)
            self.pica4 = MultiScalePiCANet(512, use_local=False, use_mid=True, use_global=True)
        elif model_type == 'no_attention':
            self.pica3 = nn.Identity()
            self.pica4 = nn.Identity()
        else:
            raise ValueError(f"Unknown model_type: {model_type}")

        self.pred_conv = nn.Conv2d(512, 1, kernel_size=1)

    def forward(self, x):
        x1, x2, x3, x4 = self.backbone(x)

        x3 = self.pica3(x3)
        if isinstance(x3, tuple):  # 兼容返回 (out, scale_weights)
            x3 = x3[0]

        x4 = self.pica4(x4)
        if isinstance(x4, tuple):
            x4 = x4[0]

        pred = self.pred_conv(x4)
        pred = F.interpolate(pred, size=(224, 224), mode='bilinear', align_corners=False)
        return pred

